import torch
from torch import nn
from prob import * 

torch.set_default_dtype(torch.float64)
torch.set_default_tensor_type(torch.DoubleTensor)


def gen_nn():
    class CDiNN(nn.Module):

        def __init__(self):
            super(CDiNN, self).__init__()

            self.act = nn.ReLU()

            self.input_layer1 = nn.Linear(DIM, D_H)
            self.input_layer1.weight.data = torch.abs(self.input_layer1.weight.data)
            self.input_layer_act1 = self.act

            self.input_layer2 = nn.Linear(DIM, D_H)
            self.input_layer2.weight.data = torch.abs(self.input_layer2.weight.data)
            self.input_layer_act2 = self.act

            self.hidden_layers1 = nn.ModuleList()
            for _ in range(N_H - 1):
                self.hidden_layer1 = nn.Linear(D_H, D_H)
                self.hidden_layer1.weight.data = torch.abs(self.hidden_layer1.weight.data)
                self.hidden_layers1.append(self.hidden_layer1)
                self.hidden_layers1.append(self.act)

            self.hidden_layers2 = nn.ModuleList()
            for _ in range(N_H - 1):
                self.hidden_layer2 = nn.Linear(D_H, D_H)
                self.hidden_layer2.weight.data = torch.abs(self.hidden_layer2.weight.data)
                self.hidden_layers2.append(self.hidden_layer2)
                self.hidden_layers2.append(self.act)

            self.output_layer_linear_prim1 = nn.Linear(D_H, 1)
            self.output_layer_linear_prim1.weight.data = torch.abs(self.output_layer_linear_prim1.weight.data)

            self.output_layer_linear_prim2 = nn.Linear(D_H, 1)
            self.output_layer_linear_prim2.weight.data = torch.abs(self.output_layer_linear_prim2.weight.data)

        def forward(self, my_input):
            x = self.input_layer1(my_input)
            x = self.input_layer_act1(x)
            for layer in self.hidden_layers1:
                x = layer(x)
            x1 = self.output_layer_linear_prim1(x)

            xx = self.input_layer2(my_input)
            xx = self.input_layer_act2(xx)
            for layer in self.hidden_layers2:
                xx = layer(xx)
            x2 = self.output_layer_linear_prim2(xx)

            return x1 - x2
    model = CDiNN()
    return model
